import mxnet as mx
from mxnet import recordio
import os
from torchvision import transforms as T
from torch.utils.data import Dataset,DataLoader
import torch
from tqdm import tqdm
import numpy as np
import cv2
from PIL import Image
# WEBface 数据集加载 : 后缀设置为rec
class WEBFACE_DATASET_REC(Dataset):
    def __init__(self,num,root,phase = 'train',input_shape = (1,112,112)):
        super(WEBFACE_DATASET_REC, self).__init__()
        #封装rec
        self.imgrec = recordio.MXIndexedRecordIO(fr"{root}/train.idx", fr"{root}/train.rec",'r')
        self.len = num
        self.phase = phase
        self.input_shape = input_shape

        if self.phase == 'train':
            self.transforms = T.Compose([
                # 转灰度图
                # 数据增样
                T.RandomCrop(self.input_shape[1:]),#裁剪成随机大小，再还原成原图尺寸，增加噪声
                T.ColorJitter(brightness=0.125, contrast=0.125, saturation=0.125),# 色彩增强，模拟在不同条件下的光照 会从均匀分布中随机采样来进行数据增强
                T.RandomHorizontalFlip(),#随机水平翻转,保证不同位置的脸部
                T.ToTensor(),
                T.Normalize((0.5), (0.5))
            ])
        else:
            self.transforms = T.Compose([
                T.CenterCrop(self.input_shape[1:]),
                T.ToTensor(),
                T.Normalize((0.5), (0.5))
            ])

    def __len__(self):
        return self.len
    def __getitem__(self, item):

        header, s = recordio.unpack(self.imgrec.read_idx(item + 1))
        img = mx.image.imdecode(s).asnumpy()
        # img = np.concatenate([img[...,[2]],img[...,[1]],img[...,[0]]],axis=2)

        img = Image.fromarray(img)
        img = img.convert('L')
        img = self.transforms(img)
        label = int(header.label)
        return img,torch.tensor(label,dtype=torch.float32)

def calc_pic_num(addr):

    num = 0
    for folder in os.listdir(addr):
        for fliename in os.listdir(f"{addr}/{folder}"):
            num +=1
    print(num)

# 计算1000个人有多少张照片
def calc_class_pic_num(addr,classnum = 1000):
    #取出前1000类

    num = 0
    allfolders  = []
    rangefolders = []
    for folder in os.listdir(addr):
        allfolders.append(folder)
        rangefolders.append(float(folder))
        # print(folder)
    # print(num)
    indexs = sorted(range(len(rangefolders)), key=lambda k: rangefolders[k])
    # print(sorted(range(len(rangefolders)), key=lambda k: rangefolders[k]))
    for i in range(classnum):
        folder_dir = fr"{addr}/{allfolders[indexs[i]]}"
        print(f"enter {folder_dir}")
        for pic in os.listdir(folder_dir):
            num+=1

    print("picture number",num)

if __name__ == '__main__':
    # calc_pic_num(r"E:\faces_webface_112x112\img")
    # calc_class_pic_num(r"E:\faces_webface_112x112\img",2000)
    data = WEBFACE_DATASET_REC(82425, r"e:\faces_webface_112x112")
    loader = DataLoader(data,batch_size=256,shuffle=True,num_workers=1)
    # print(data[1][0].shape)

    for i,(img,label) in enumerate(tqdm(loader)):
        try :
            print(label.tolist())
        except:
            print(i)
